#include "typechecker.hpp"
#include "symtable.hpp"
#include "ast.hpp"

#include <cassert>
#include <iostream>
using namespace std;

bool TypeChecker::has_error() {
    return _has_error;
}

void TypeChecker::visit(Node* n) {
    n->accept(this);
}

void TypeChecker::visit(PrintNode* n) {
    visit(n->value());
}

void TypeChecker::visit(ReturnNode* n) {
    visit(n->value());
    if (!_sym_table->is_func()) {
        cout << "Can not use return outside of function" << endl;
        _has_error = true;
        return;
    }
    _sym_table->set_ret_type(_type);
}

void TypeChecker::visit(BinaryOp* op) {
    _type = nullptr;
    visit(op->left());
    Type* left_type = _type;

    _type = nullptr;
    visit(op->right());
    Type* right_type = _type;

    if (!left_type->equals(right_type)) {
        cerr << "Unsupported binary operator between " << left_type->to_string()
            << " and " << right_type->to_string() << endl;
        _has_error = true;
        return;
    }

    switch(op->op_type()) {
    case AST_OP_ADD:
        if (!left_type->support_operator(OPERATOR_ADD)) {
            cerr << "Unsupported operator '+' for type " << left_type->to_string() << endl;
            _has_error = true;
        }
        break;
    case AST_OP_SUB:
        break;
    case AST_OP_MUL:
        break;
    case AST_OP_DIV:
        break;
    case AST_OP_BIT_AND:
        break;
    case AST_OP_BIT_OR:
        break;
    case AST_OP_BIT_XOR:
        break;
    case AST_OP_LOG_OR:
        break;
    case AST_OP_LOG_AND:
        break;
    case AST_OP_LEFT_SHIFT:
        break;
    case AST_OP_RIGHT_SHIFT:
        break;
    default:
        printf("Unknown binary op %d\n", op->op_type());
        return;
    }
}

void TypeChecker::visit(ListNode* n) {
    for (auto it = n->node_list()->begin(); 
        it != n->node_list()->end(); it++) {
        visit(*it);
    }
}

void TypeChecker::visit(ConstInt* n) {
    _type = IntType::get_instance();
}

void TypeChecker::visit(BuiltinConst* n) {
    switch (n->_builtin_id) {
    case BuiltinConst::UNIT:
        _type = UnitType::get_instance();
        return;
    case BuiltinConst::TRUE:
    case BuiltinConst::FALSE:
        _type = BoolType::get_instance();
        return;
    default:
        cout << "unknown const: " << n->_builtin_id;
    }
}

void TypeChecker::visit(ConstString* n) {
    _type = StringType::get_instance();
}

void TypeChecker::visit(ConstArray* n) {
    visit(n->arr->get(0));
    _type = new ArrayType(_type, n->arr->size());
}

void TypeChecker::visit(IfNode* n) {
    visit(n->cond());
    visit(n->then_part());

    if (n->else_part() != NULL) {
        visit(n->else_part());
    }
}

void TypeChecker::visit(WhileNode* n) {
    visit(n->cond());
    visit(n->body());
}

void TypeChecker::visit(CmpNode* n) {
    visit(n->left());
    visit(n->right());
    if (n->cmp_op() == CmpNode::CMP_LT) {
    }

    _type = BoolType::get_instance();
}

void TypeChecker::visit(LogicNotNode* n) {
    visit(n->value());
    _type = BoolType::get_instance();
}

void TypeChecker::visit(AssignNode* n) {
    _type = nullptr;
    visit(n->left());
    Type* left_type = _type;

    _type = nullptr;
    visit(n->right());
    Type* right_type = _type;

    if (!left_type->equals(right_type)) {
        cerr << "Incompatible assigning from " << right_type->to_string()
            << " to " << left_type->to_string() << endl;
        _has_error = true;
    }
}

void TypeChecker::visit(VarNode* n) {
    int index = _sym_table->get_name_index(n->name());
    if (index >= 0) {
        _type = _sym_table->get_name_type(n->name());
        return;
    }

    index = _sym_table->get_var_name_index(n->name());
    if (index >= 0) {
        _type = _sym_table->get_var_name_type(n->name());
        return;
    }

    index = _sym_table->get_cell_var_index(n->name());
    if (index >= 0) {
        _type = _sym_table->get_cell_var_type(n->name());
        return;
    }

    index = _sym_table->get_free_var_index(n->name());
    if (index >= 0) {
        _type = _sym_table->get_free_var_type(n->name());
        return;
    }

    // find name along the sym table chain.
    SymbolTable* pst = _sym_table->next;
    while (pst != nullptr) {
        index = pst->get_name_index(n->name());
        if (index >= 0) {
            _type = pst->get_name_type(n->name());
            _sym_table->add_name(n->name(), _type);
            break;
        }

        if (!pst->is_func()) {
            pst = pst->next;
            continue;
        }

        // If pst is func, try to find variable in varnames.
        index = pst->get_var_name_index(n->name());
        if (index >= 0) {
            _type = pst->get_var_name_type(n->name());

            /* For those functions that variable is not defined in,
             * this variable is free, else it is cell.
             */
            SymbolTable* begin = _sym_table;
            while (begin != pst) {
                if (begin->get_free_var_index(n->name()) < 0) {
                    begin->add_free_var(n->name(), _type);
                }
                begin = begin->next;
            }

            if (pst->get_cell_var_index(n->name()) < 0) {
                pst->add_cell_var(n->name());
            }
            break;
        }
        pst = pst->next;
    }

    if (index == -1) {
        cerr << "undeclared variable '" << n->name() << "'." << endl;
        _has_error = true;
    }
}

void TypeChecker::visit(CallNode* n) {
    visit(n->callee());

    if (_type->is_array()) {
        ArrayType* at = dynamic_cast<ArrayType*>(_type);
        delete n->callee();
        n->set_callee(new VarNode("make_array"));

        // revisit function to fill symtable.
        Type* old_type = _type;

        visit(n->callee());
        n->args()->add(new ConstInt(at->size()));
        for (auto it = n->args()->node_list()->begin();
            it != n->args()->node_list()->end(); it++) {
            visit(*it);
        }

        _type = old_type;
        return;
    }

    Type* callee_type = _type;

    if (n->args()->size() == 0) {
        _type = callee_type->try_apply(UnitType::get_instance());
        return;
    }

    // apply all arguments one by one.
    for (auto it = n->args()->node_list()->begin();
        it != n->args()->node_list()->end(); it++) {
        visit(*it);

        callee_type = callee_type->try_apply(_type);
        if (!callee_type) {
            cout << "Call function with wrong arguments." << endl;
            _has_error = true;
            break;
        }
    }

    _type = callee_type;
}

void TypeChecker::visit(FuncDefNode* n) {
    if (n->has_generic_args()) {
        _type = GenericFuncType::get_instance();
        return;
    }

    _sym_table = n->sym_table();

    visit(n->args());
    visit(n->body());

    if (_sym_table->is_func()) {
        _type = _sym_table->get_var_name_type(n->name());
    }
    else {
        _type = _sym_table->get_name_type(n->name());
    }

    _sym_table = _sym_table->next;
}

void TypeChecker::instantiate_generic_func(FuncDefNode* n) {
    /* 清理上一次实例化时的状态 */
    if (n->has_generic_args()) {
        if (n->sym_table())
            delete n->sym_table();
    }

    n->set_sym_table(new SymbolTable(_sym_table, true));
    _sym_table = n->sym_table();
    _sym_table->set_arg_count(n->args()->size());

    _type_table = n->type_table();

    Type* func_type = build_func_type(n);

    visit(n->args());
    visit(n->body());

    /* Return type got from return statement. 
     * It should be same as declared type of function.
     * ret_type equals declared return type, 
     * this is made sure by ReturnNode. 
     */
    if (!n->ret_type()) {
        visit(n->ret_type());
        _sym_table->set_ret_type(_type);
    }

    _type = func_type;

    /* Restore to parent's scope */
    _type_table = _type_table->next();
    _sym_table = _sym_table->next;
}

void TypeChecker::visit(SubscrNode* n) {
    visit(n->name());
    if (_type->is_generic_func()) {
        ListNode* ln = dynamic_cast<ListNode*>(n->index());
        assert(ln != nullptr);

        VarNode* name_node = dynamic_cast<VarNode*>(n->name());
        FuncDefNode* func_def = dynamic_cast<FuncDefNode*>(_sym_table->get_subtree_by_name(name_node->name()));
        assert(func_def != nullptr);

        int i = 0;
        for (auto it = ln->begin(); it != ln->end(); it++) {
            visit(*it);
            func_def->type_table()->set_generic_name(i++, _type);
        }

        instantiate_generic_func(func_def);
        n->set_generic_inst(1);
        return;
    }

    visit(n->index());
    assert(_type == IntType::get_instance());
    ArrayType* t = dynamic_cast<ArrayType*>(_type);
    if (t != nullptr) {
        _type = t->elm_type();
    }
}

Type* TypeChecker::get_param_type(ParamDefNode* n) {
    if (n->def_type()) {
        visit((Node*)n->def_type());
        return _type;    
    }

    return nullptr;;
}

Type* TypeChecker::build_func_type(FuncDefNode* n) {
    ListNode* args = n->args();

    visit(n->ret_type());
    Type* func_type = _type;

    if (args->size() == 0) {
        return new ArrowType(UnitType::get_instance(), func_type);
    }

    for (auto it = args->node_list()->rbegin(); 
        it != args->node_list()->rend(); it++) {
        Type* t = get_param_type(dynamic_cast<ParamDefNode*>(*it));
        func_type = new ArrowType(t, func_type);
    }

    return func_type;
}

void TypeChecker::visit(ParamDefNode* n) {
    if (_sym_table->get_var_name_index(n->name()) != -1) {
        cerr << "redefination for variable '" << n->name() << "'." << endl;
        _has_error = true;
        return;
    }

    Type* declare_type = get_param_type(n);
    _sym_table->add_var_name(n->name(), declare_type);
}

/*
 * All local variables are put in varnames.
 */
template<bool is_func>
void TypeChecker::check_name(const string& name, Type* t) {
    int index = -1;
    index = _sym_table->get_var_name_index(name);
    if (index != -1) {
        cerr << "redefination for variable '" << name << "'." << endl;
        _has_error = true;
        return;
    }
    _sym_table->add_var_name(name, t);
}

template<>
void TypeChecker::check_name<false>(const string& name, Type* t) {
    int index = -1;
    index = _sym_table->get_name_index(name);
    if (index != -1) {
        cerr << "redefination for variable '" << name << "'." << endl;
        _has_error = true;
        return;
    }
    _sym_table->add_name(name, t);
}

void TypeChecker::visit(VarDefNode* n) {
    Type* declare_type = nullptr;
    Type* value_type = nullptr;

    if (n->def_type()) {
        _type = nullptr;
        visit(n->def_type());
        declare_type = _type;    
    }

    if (n->init_value()) {
        _type = nullptr;
        visit(n->init_value());
        value_type = _type;
    }

    // check type rules.
    if (declare_type && value_type) {
        if (!declare_type->equals(value_type)) {
            cerr << "types conflict for " << n->name()  << endl;
            _has_error = true;
        }
    }
    else if (value_type) {
        declare_type = value_type;
    }
    else {
        cerr << "declare variable '" << n->name() << "' without type." << endl;
        _has_error = true;
    }

    if (_sym_table->is_func()) {
        check_name<true>(n->name(), declare_type);
    }
    else {
        check_name<false>(n->name(), declare_type);
    }
}

void TypeChecker::visit(TypeVarNode* n) {
    _type = _type_table->get_expanded_type(n->name());
    if (_type == nullptr) {
        cerr << "Unknown type: '" << n->name() << "'." << endl;
        _has_error = true;
    }
}

void TypeChecker::visit(TypeVarDefNode* n) {
    visit(n->value());
    _type_table->add(n->name(), _type);
}

void TypeChecker::visit(TypeFuncDefNode* n) {
    TypeFunction* tf = new TypeFunction(n->body());
    for (auto it = n->args()->node_list()->begin(); 
        it != n->args()->node_list()->end(); it++) {
        TypeVarNode* tvn = dynamic_cast<TypeVarNode*>(*it);
        tf->args().push_back(tvn->name());
    }
    _type = tf;
    _type_table->add(n->name(), _type);
}

void TypeChecker::visit(TypeCallNode* n) {
    visit(n->callee());
    TypeFunction* tf = dynamic_cast<TypeFunction*>(_type);
    if (tf == nullptr) {
        cout << "Can not instantiate an generic type." << endl;
        return;
    }

    int i = 0;
    for (auto it = n->args()->node_list()->begin(); 
        it != n->args()->node_list()->end(); it++) {
        visit(*it);
        _type_table->add(tf->args()[i++], _type);
    }

    visit(tf->body());
}

void TypeChecker::visit(TypeConstInt* n) {
    _type = new TypeDependentConst(n->value);
}

void TypeChecker::visit(ArrowTypeNode* n) {
    visit((Node*)n->src());
    Type* left = _type;
    visit((Node*)n->dst());
    Type* right = _type;
    _type = new ArrowType(left, right);
}

/*
 * 检查函数定义处的类型，将函数名称放入symtable。
 */
void TypeChecker::fill_sym_table(FuncDefNode* n) {
    /* Recursion func needs func name. 
     * Add name to name table first.
     */
    Type* func_type = nullptr;
    if (n->has_generic_args()) {
        func_type = GenericFuncType::get_instance();
    }
    else {
        func_type = build_func_type(n);
    }

    // nested func declaration.
    if (_sym_table->is_func()) {
        check_name<true>(n->name(), func_type);
    }
    else {
        check_name<false>(n->name(), func_type);
    }

    /* 如果是泛型函数，就先不切换sym table，因为这里还不能检查函数的定义是否正确 */
    if (n->has_generic_args()) {
        _sym_table->set_subtree(n->name(), n);

        n->set_type_table(new TypeTable(_type_table));
        _type_table = n->type_table();
        for (auto it = n->generic_args()->begin(); 
            it != n->generic_args()->end(); it++) {
            _type_table->add_generic_name(*it);
        }
    }

    /* 创建函数所对应的符号表，记录返回类型 */
    n->set_sym_table(new SymbolTable(_sym_table, true));
    _sym_table = n->sym_table();
    _sym_table->set_arg_count(n->args()->size());

    /* Return type got from return statement. 
     * It should be same as declared type of function.
     * ret_type equals declared return type, 
     * this is made sure by ReturnNode. 
     */
    if (n->ret_type()) {
        visit(n->ret_type());
        _sym_table->set_ret_type(_type);
    }

    /* Restore to parent's scope */
    _type_table = _type_table->next();
    _sym_table = _sym_table->next;
    _type = func_type;
}

void TypeChecker::visit(LambdaDef* n) {
    /* Create new symtable for function. */
    /* In function scope. */
    n->set_sym_table(new SymbolTable(_sym_table, true));
    _sym_table = n->sym_table();
    _sym_table->set_arg_count(1);

    visit(n->param());
    Type* param_type = _type;

    visit(n->body());
    _type = new ArrowType(param_type, _sym_table->ret_type());

    /* Restore to parent's scope */
    _sym_table = _sym_table->next;
}

/* pseudo nodes */
void TypeChecker::visit(ArrayTypeNode* n) {
    visit(n->elm_type());
    Type* elm_type = _type;
    visit(n->size());
    _type = new ArrayType(elm_type, dynamic_cast<TypeDependentConst*>(_type)->value);
}

void Evaluator::visit(Node* n) {
    n->accept(this);
}

void Evaluator::visit(TypeConstInt* n) {
    _value.cvt = CVT_INT;
    _value.value.int_value = n->value;
}

